package task

import (
	"errors"
	"fmt"
	"sync"
)

type Task struct {
	wg      sync.WaitGroup
	lock    sync.Mutex
	works   []*Work
	mapping map[string]int
}

func (t *Task) exec(work *Work) {
	defer t.wg.Done()
	defer close(work.finish)

	d, e := work.fn(t)
	if e == nil {
		work.Data = Data{
			Result: true,
			Err:    nil,
			Data:   d,
		}
	} else {
		work.Data = Data{
			Result: false,
			Err:    e,
			Data:   nil,
		}
	}
}

func (t *Task) Get(name string) (w Work, e error) {
	t.lock.Lock()
	defer t.lock.Unlock()

	key, ok := t.mapping[name]
	if !ok {
		e = errors.New(fmt.Sprintf("名为%s的work不存在", name))
		return
	}

	return *t.works[key], nil
}

func (t *Task) Add(name string, fn WorkFunHandler) *Task {
	t.lock.Lock()
	defer t.lock.Unlock()

	key := len(t.works)
	t.works = append(t.works, &Work{
		fn:     fn,
		Name:   name,
		finish: make(chan struct{}),
	})
	t.mapping[name] = key

	return t
}

func (t *Task) Wait() *Task {
	for _, w := range t.works {
		t.wg.Add(1)
		go t.exec(w)
	}
	t.wg.Wait()
	return t
}

func (t *Task) Err() error {
	for _, w := range t.works {
		if !w.Data.Result {
			return w.Data.Err
		}
	}
	return nil
}

func (t *Task) Result() map[string]Data {
	r := make(map[string]Data)
	for _, w := range t.works {
		r[w.Name] = w.Data
	}

	return r
}

func NewTask() *Task {
	return &Task{
		wg:      sync.WaitGroup{},
		lock:    sync.Mutex{},
		works:   make([]*Work, 0),
		mapping: make(map[string]int),
	}
}
